"""
消息发送限速中间件
在bot连接时hook send方法，添加限速逻辑
"""
from typing import TYPE_CHECKING, Any, Optional
from collections import defaultdict
import time
from nonebot import get_driver
from nonebot.adapters.onebot.v11 import Bot, GroupMessageEvent
from src.common.rate_limiter import GroupRateLimiter

if TYPE_CHECKING:
    from nonebot.adapters.onebot.v11 import Message, MessageSegment

# 配置参数（可以从配置文件读取）
RATE_LIMIT_CONFIG = {
    "enabled": True,  # 是否启用限速
    "capacity": 5,    # 桶容量（允许突发5条）
    "refill_rate": 2.0,  # 补充速率（每秒2条）
    "per_group": False,  # True=按群限速，False=全局限速
    "emergency_stop": False,  # 紧急停止开关
    "auto_stop_threshold": 20,  # 单个群连续发送超过此数量时自动紧急停止
    "auto_stop_time_window": 60,  # 时间窗口（秒），在此时间内连续发送才会触发自动停止
}

# 创建限速器实例
limiter = GroupRateLimiter(
    capacity=RATE_LIMIT_CONFIG["capacity"],
    refill_rate=RATE_LIMIT_CONFIG["refill_rate"]
)

# 跟踪每个群的消息计数和时间
group_message_tracker: dict[int, dict] = defaultdict(lambda: {"count": 0, "first_msg_time": None})


async def check_and_update_group_message_count(group_id: Optional[int], bot: Bot, event: Optional[GroupMessageEvent] = None):
    """
    检查并更新群消息计数，如果超过阈值则自动触发紧急停止
    
    Args:
        group_id: 群组ID
        bot: Bot实例（用于发送通知消息）
        event: 事件实例（用于发送通知消息，可选）
    """
    if group_id is None:
        return  # 不是群消息，不处理
    
    now = time.time()
    tracker = group_message_tracker[group_id]
    
    # 检查时间窗口，如果超出窗口则重置计数
    if tracker["first_msg_time"] is None:
        # 第一次发送消息，初始化
        tracker["first_msg_time"] = now
        tracker["count"] = 1
    else:
        elapsed = now - tracker["first_msg_time"]
        if elapsed > RATE_LIMIT_CONFIG["auto_stop_time_window"]:
            # 超出时间窗口，重置计数
            tracker["first_msg_time"] = now
            tracker["count"] = 1
        else:
            # 在时间窗口内，增加计数
            tracker["count"] += 1
    
    print(f"[限速器] 群 {group_id} 连续消息计数: {tracker['count']}/{RATE_LIMIT_CONFIG['auto_stop_threshold']}")
    
    # 检查是否超过阈值
    if tracker["count"] >= RATE_LIMIT_CONFIG["auto_stop_threshold"]:
        print(f"[限速器] ⚠️ 群 {group_id} 在 {RATE_LIMIT_CONFIG['auto_stop_time_window']} 秒内发送了 {tracker['count']} 条消息，超过阈值 {RATE_LIMIT_CONFIG['auto_stop_threshold']}，自动触发紧急停止！")
        
        # 自动触发紧急停止
        RATE_LIMIT_CONFIG["emergency_stop"] = True
        
        # 尝试发送通知消息
        try:
            original_send = getattr(bot, '_original_send', None)
            notification_msg = f"🛑 自动紧急停止已触发！\n检测到群内 {RATE_LIMIT_CONFIG['auto_stop_time_window']} 秒内连续发送了 {tracker['count']} 条消息（阈值: {RATE_LIMIT_CONFIG['auto_stop_threshold']} 条）。\n使用「恢复限速」命令可以恢复。"
            
            # 优先使用event发送（通过bot.send）
            if original_send and event:
                await original_send(event=event, message=notification_msg)
            # 如果没有event，尝试使用call_api直接发送群消息
            elif getattr(bot, '_original_call_api', None):
                await bot._original_call_api('send_group_msg', group_id=group_id, message=notification_msg)
        except Exception as e:
            print(f"[限速器] ⚠️ 无法发送自动紧急停止通知: {e}")


def reset_group_message_count(group_id: Optional[int]):
    """重置群消息计数"""
    if group_id is not None:
        group_message_tracker[group_id] = {"count": 0, "first_msg_time": None}
        print(f"[限速器] 群 {group_id} 的消息计数已重置")


def setup_rate_limiter_for_bot(bot: Bot):
    """
    Bot连接时的钩子函数
    替换bot.send方法，添加限速逻辑
    """
    print("[限速器] 正在初始化消息发送限速中间件...")

    # 保存原始的send方法（存储到bot实例上，方便后续访问）
    original_send = bot.send
    bot._original_send = original_send  # 保存到bot实例上

    async def rate_limited_send(*args, **kwargs):
        """
        带限速的消息发送方法

        这个方法会拦截所有 bot.send() 调用
        使用 *args, **kwargs 来兼容不同的调用方式
        
        注意：限速和计数都在 call_api 中统一处理，因为 send 内部会调用 call_api
        这样可以避免重复限速和计数
        """
        print(f"[限速器] 🔍 拦截到 bot.send 调用")
        
        # 只检查紧急停止（限速和计数都在 call_api 中统一处理）
        if RATE_LIMIT_CONFIG["emergency_stop"]:
            print("[限速器] ⚠️ 紧急停止已启用，消息发送被阻止")
            raise RuntimeError("紧急停止：消息发送已被阻止")

        # 直接调用原始send方法（限速和计数会在 call_api 中处理）
        result = await original_send(*args, **kwargs)
        print(f"[限速器] ✅ bot.send 调用完成")
        
        return result

    # 同时hook call_api方法（NoneBot可能通过call_api发送消息）
    original_call_api = bot.call_api
    bot._original_call_api = original_call_api
    
    async def rate_limited_call_api(api: str, **data):
        """
        带限速的API调用方法
        拦截send_msg相关的API调用
        """
        # 只对消息发送相关的API进行限速
        message_apis = ['send_msg', 'send_private_msg', 'send_group_msg']
        
        if api in message_apis:
            print(f"[限速器] 🔍 拦截到 call_api: {api}")
            
            # 检查是否启用限速
            if not RATE_LIMIT_CONFIG["enabled"]:
                return await original_call_api(api, **data)

            # 检查紧急停止
            if RATE_LIMIT_CONFIG["emergency_stop"]:
                print("[限速器] ⚠️ 紧急停止已启用，API调用被阻止")
                raise RuntimeError("紧急停止：消息发送已被阻止")

            # 提取群组ID（send_msg API可能通过message_type判断，send_group_msg直接有group_id）
            group_id = None
            if api == 'send_group_msg':
                group_id = data.get('group_id')
            elif api == 'send_msg':
                # send_msg API可能需要根据message_type判断
                if data.get('message_type') == 'group' and 'group_id' in data:
                    group_id = data['group_id']
            
            if group_id:
                print(f"[限速器] 检测到群消息，群号: {group_id}")

            # 获取令牌
            print(f"[限速器] ⏳ 等待令牌... (模式: {'按群' if RATE_LIMIT_CONFIG['per_group'] else '全局'})")
            try:
                if RATE_LIMIT_CONFIG["per_group"]:
                    await limiter.acquire(group_id)
                else:
                    await limiter.acquire(None)
                print(f"[限速器] ✅ 获得令牌，准备调用API")
            except Exception as e:
                print(f"[限速器] ❌ 获取令牌时出错: {e}")
                raise

            result = await original_call_api(api, **data)
            print(f"[限速器] ✅ API调用完成")
            
            # 只在 call_api 中计数（这是真正发送消息的地方）
            # 这样可以避免 bot.send() 和 call_api() 重复计数的问题
            if group_id is not None:
                # 注意：call_api中没有event，所以传None，函数会使用call_api发送通知
                await check_and_update_group_message_count(group_id, bot, None)
            
            return result
        else:
            # 非消息发送API，直接调用
            return await original_call_api(api, **data)

    # 替换bot的方法
    bot.send = rate_limited_send
    bot.call_api = rate_limited_call_api

    print(f"[限速器] ✅ 限速中间件已启用")
    print(f"  - 桶容量: {RATE_LIMIT_CONFIG['capacity']} 条")
    print(f"  - 限速速率: {RATE_LIMIT_CONFIG['refill_rate']} 条/秒")
    print(f"  - 限速模式: {'按群限速' if RATE_LIMIT_CONFIG['per_group'] else '全局限速'}")


# 注册bot连接事件监听器
driver = get_driver()


@driver.on_bot_connect
async def on_bot_connect(bot: Bot):
    """Bot连接时自动调用"""
    setup_rate_limiter_for_bot(bot)


# 控制命令
from nonebot import on_command

print("[限速器] 正在注册控制命令...")

enable_rate_limit_cmd = on_command("启用限速", priority=1)
disable_rate_limit_cmd = on_command("禁用限速", priority=1)
emergency_stop_cmd = on_command("紧急停止", priority=1)
resume_rate_limit_cmd = on_command("恢复限速", priority=1)  # 恢复限速（取消紧急停止）
rate_limit_status_cmd = on_command("限速状态", priority=1)
switch_global_mode_cmd = on_command("全局限速", priority=1)
switch_group_mode_cmd = on_command("按群限速", priority=1)

print("[限速器] ✅ 控制命令已注册: 启用限速, 禁用限速, 紧急停止, 恢复限速, 限速状态, 全局限速, 按群限速")


@enable_rate_limit_cmd.handle()
async def enable_rate_limit(bot: Bot, event: GroupMessageEvent):
    """启用限速"""
    RATE_LIMIT_CONFIG["enabled"] = True
    RATE_LIMIT_CONFIG["emergency_stop"] = False
    await enable_rate_limit_cmd.finish("✅ 限速器已启用")


@disable_rate_limit_cmd.handle()
async def disable_rate_limit(bot: Bot, event: GroupMessageEvent):
    """禁用限速"""
    RATE_LIMIT_CONFIG["enabled"] = False
    await disable_rate_limit_cmd.finish("✅ 限速器已禁用")


@emergency_stop_cmd.handle()
async def emergency_stop_handler(bot: Bot, event: GroupMessageEvent):
    """紧急停止所有消息发送"""
    # 先发送确认消息（在设置紧急停止之前）
    try:
        # 使用保存的原始send方法绕过限速器检查
        original_send = getattr(bot, '_original_send', None)
        if original_send:
            await original_send(event=event, message="🛑 紧急停止已启用！\n所有后续消息发送将被阻止。\n使用「恢复限速」命令可以恢复。")
        else:
            # 如果找不到原始方法，直接使用bot.send（在紧急停止设置前）
            await bot.send(event=event, message="🛑 紧急停止已启用！\n所有后续消息发送将被阻止。\n使用「恢复限速」命令可以恢复。")
    except Exception as e:
        # 如果发送失败，至少打印日志
        print(f"[限速器] ⚠️ 无法发送紧急停止确认消息: {e}")
    
    # 设置紧急停止（在所有确认消息发送后再设置）
    RATE_LIMIT_CONFIG["emergency_stop"] = True
    RATE_LIMIT_CONFIG["enabled"] = True  # 保持限速器启用状态
    
    # 不再调用finish，因为已经手动发送了消息
    print(f"[限速器] ⚠️ 紧急停止已启用（群号: {event.group_id}）")


@resume_rate_limit_cmd.handle()
async def resume_rate_limit(bot: Bot, event: GroupMessageEvent):
    """恢复限速（取消紧急停止）"""
    # 先取消紧急停止
    RATE_LIMIT_CONFIG["emergency_stop"] = False
    RATE_LIMIT_CONFIG["enabled"] = True
    
    # 重置当前群的消息计数
    reset_group_message_count(event.group_id)
    
    # 发送恢复消息（现在可以正常发送了）
    try:
        original_send = getattr(bot, '_original_send', None)
        if original_send:
            await original_send(event=event, message="✅ 限速器已恢复！\n紧急停止已取消，消息计数已重置，消息发送恢复正常。")
        else:
            await bot.send(event=event, message="✅ 限速器已恢复！\n紧急停止已取消，消息计数已重置，消息发送恢复正常。")
    except Exception as e:
        print(f"[限速器] ⚠️ 无法发送恢复确认消息: {e}")
    
    print(f"[限速器] ✅ 紧急停止已取消，限速器恢复正常（群号: {event.group_id}）")


@rate_limit_status_cmd.handle()
async def rate_limit_status(bot: Bot, event: GroupMessageEvent):
    """查看限速状态"""
    print(f"[限速器] 收到限速状态查询命令，群号: {event.group_id}")
    status = "启用" if RATE_LIMIT_CONFIG["enabled"] else "禁用"
    stop = "是" if RATE_LIMIT_CONFIG["emergency_stop"] else "否"
    mode = "按群限速" if RATE_LIMIT_CONFIG["per_group"] else "全局限速"

    # 显示当前群的计数信息
    current_count = group_message_tracker.get(event.group_id, {}).get("count", 0)
    
    msg = f"""📊 限速器状态
状态: {status}
紧急停止: {stop}
限速模式: {mode}
桶容量: {RATE_LIMIT_CONFIG['capacity']} 条
限速速率: {RATE_LIMIT_CONFIG['refill_rate']} 条/秒
自动停止阈值: {RATE_LIMIT_CONFIG['auto_stop_threshold']} 条/{RATE_LIMIT_CONFIG['auto_stop_time_window']} 秒
当前群计数: {current_count}/{RATE_LIMIT_CONFIG['auto_stop_threshold']}"""

    await rate_limit_status_cmd.finish(msg)


@switch_global_mode_cmd.handle()
async def switch_global_mode(bot: Bot, event: GroupMessageEvent):
    """切换到全局限速模式"""
    RATE_LIMIT_CONFIG["per_group"] = False
    await switch_global_mode_cmd.finish("✅ 已切换到全局限速模式\n所有群共享同一个限速桶")


@switch_group_mode_cmd.handle()
async def switch_group_mode(bot: Bot, event: GroupMessageEvent):
    """切换到按群限速模式"""
    RATE_LIMIT_CONFIG["per_group"] = True
    await switch_group_mode_cmd.finish("✅ 已切换到按群限速模式\n每个群使用独立的限速桶")

